#include <profo/SimpleHP2DRectangularAlgorithm.h>
#include <profo/profo.h>
#include <profo/util.h>
#include <iostream>

using namespace profo;

/**
 * The one and only SimpleHP2DRectangularAlgorithm constructor.
 * @param arg A reference to the protein object that is going to be used in
 *            the simulation process.
 */
SimpleHP2DRectangularAlgorithm::SimpleHP2DRectangularAlgorithm(Protein& arg):
  IAlgorithm(arg),
  _proteinRef(arg),
  _authors("Dejan Lekic"),
  _description("Simple Local Search Algorithm")
{
  _proteinRef.authors(_authors);
  _proteinRef.description(_description);
  _proteinRef.energy(-2.0); // TODO: this should be calculated here!
} // SimpleHP2DRectangularAlgorithm constructor

/**
 * Given the move argMove calculate new X and Y coordinate.
 * NOTE:
 * Resulting X and Y coordinates will be stored in argX and argY arguments upon
 * return.
 *
 * @param argId ID of an amino acid that is about to move in direction argMove
 * @param argMove A number from 0 to 7 which specifies which move it is.
 */
void SimpleHP2DRectangularAlgorithm::calculateXY(unsigned int argId, unsigned char argMove,
                                                 int& argX,
                                                 int& argY) {
  AminoAcid& aa = _proteinRef.get(argId);
  switch (argMove) {
    case 0: argX = aa.x + 1; break;
    case 1: argX = aa.x + 1; argY = aa.y + 1; break;
    case 2: argY = aa.y + 1; break;
    case 3: argY = aa.y + 1; argX = aa.x - 1; break;
    case 4: argX = aa.x - 1; break;
    case 5: argX = aa.x - 1; argY = aa.y - 1; break;
    case 6: argY = aa.x - 1; break;
    case 7: argX = aa.x + 1; argY = aa.y - 1; break;
  } // switch
} // calculateXY()

/**
 * This method checks validity of a currently processed Protein.
 * @param arg Stop checking further if we reach amino acid with ID = arg.
 *            If arg is zero, we will check whole Protein.
 */
bool SimpleHP2DRectangularAlgorithm::checkValidity(unsigned int arg) const {
  bool ret = true;
  int limit = (arg == 0) ? _proteinRef.size() : arg+1;
  for (int i = 0; i < limit; i++) {
    AminoAcid& first = _proteinRef.get(i);
    AminoAcid& second = _proteinRef.get(i+1);
    if (!topologicalNeighbour(first, second)) {
      return false;
    } // if
  } // false
} // checkValidity()

/**
 * One of the most vital methods - it executes a move specified in the argument.
 * NOTE: We rely a lot on correctness of the getAvailableMoves() method
 * so we must test it method thoroughly!
 *
 * Here is a description of moves
 * @param arg An input value describing which move to do.
 * 0 - move east          4 - move west
 * 1 - move north-east    5 - move south-west
 * 2 - move north         6 - move south
 * 3 - move north-west    7 - move south-east
 */
void SimpleHP2DRectangularAlgorithm::move(unsigned char arg) {
  int id = _proteinRef.selected();
  // Check again if move is valid
  unsigned char available = getAvailableMoves(id);
  if (!getBit(arg, available)) // if the given move "arg" is NOT available
    return;                    // we just exit.

  // if selected amino acid is the first amino acid, then we just do the move
  // because it is definitely a valid move.
  if (id == 0) {
    _proteinRef.move(id, arg);
    // re-calculate the energy
    _proteinRef.energy(energy()); // update the energy
    updateFreeMoves(); // we must update free moves now
    return;
  } // if

  // Any other move must check the C-location, and see if we can just move selected
  // amino acid to a new location, or we must perform a sequence of moves, like
  // it is explaned in (Lech et al., 2003, 189)
  // let's first get coordinates of the L position
  int lx = 0;
  int ly = 0;
  calculateXY(id, arg, lx, ly);
  // we will store C-coordinates in the following two variables
  int cx = lx;
  int cy = ly;
#ifdef DEBUG
      std::cout << "calling checkCorner(" << id <<","<<cx<<","<<cy <<")"<< std::endl;
#endif
  if (checkCorner(id, cx, cy)) {
    if (_moveDirectly) {
      // we can safely move currently selected amino acid and polypeptide
      // conformation will remain valid
      _proteinRef.move(id, arg);
#ifdef DEBUG
      std::cout << "direct move" << std::endl;
#endif
      return;
    } else {
      // this is the most complex case - when C location is free.
      // we must first store coordinates of the currently selected amino acid
      // because soon they will be changed.
      AminoAcid& current = _proteinRef.get(arg);
      int curx = current.x;
      int cury = current.y;
      // #1 we move selected amino acid to location L
      _proteinRef.move(id, arg);
      // #2 we move previous amino acid in the chain (id = id -1) to location C
      AminoAcid& previous = _proteinRef.get(arg-1);
      // we must store coordinates
      int prevx = previous.x;
      int prevy = previous.y;
      // do the move
#ifdef DEBUG
      std::cout << "C:" << cx << "," << cy << std::endl;
#endif
      previous.x = cx;
      previous.y = cy;
      // here we will store temporary coordinates
      int tmpx = 0;
      int tmpy = 0;
      // #3 starting from amino acid with id=id-2 we loop backwards till we reach
      // a valid conformation.
      for (int j=arg-2; ((j >= 0) && !checkValidity(arg)); j--) {
        AminoAcid& aa = _proteinRef.get(j);
        tmpx = aa.x;
        tmpy = aa.y;
        aa.x = curx;
        aa.y = cury;
        curx = prevx;
        cury = prevy;
        prevx = tmpx;
        prevy = tmpy;
      } // for
    } // else
  } // else
} // move() method

/**
 * Checks if given two AminoAcids are topological neighbors.
 * NOTE: This is implementation-specific function. Some algorithms can consider
 * diagonal neighbours too. This implementation does not count them as neighbours.
 */
bool SimpleHP2DRectangularAlgorithm::topologicalNeighbour(const AminoAcid& first, const AminoAcid& second) const {
  int fx = first.x;
  int fy = first.y;
  int sx = second.x;
  int sy = second.y;
  return (((fx == sx-1) && (fy == sy)) ||   // east
         ((fx == sx) && (fy == sy-1)) ||    // north
         ((fx == sx+1) && (fy == sy)) ||    // west
         ((fx == sx) && (fy == sy+1)));     // south
} // topologicalNeighbour() method

/**
 * Checks if there is no amino acid at location (argX,argY).
 *
 * Complexity: O(n) where n = number of amino acids in the protein chain.
 */
bool SimpleHP2DRectangularAlgorithm::isFree(int argX, int argY) {
  int tmpX, tmpY;
  bool ret = true;
  for (int i=0; i<_proteinRef.size(); i++) {
    tmpX = _proteinRef.get(i).x;
    tmpY = _proteinRef.get(i).y;
    // we have found an amino acid on given location, return false
    if ((tmpX == argX) && (tmpY == argY)) {
      return false;
    } // if
  } // for
} // isFree() method

/**
 * This method checks whether corner location (location "C" in the publication)
 * is either free, or an amino acid with ID = argId-1 is on that position. If it is
 * so, method returns true, in other words - residue can move to location
 * (argX, argY).
 * NOTE: This method also stores information whether corner is free or not
 * in the _cornerFree member variable, so we do not have to use isFree() again
 * when we perform the move. Secondly, _moveDirectly is set as well.
 *
 * WARNING: This method will _OVERWRITE_ argX and argY variables! If you want to
 * use old valus, store arguments somewhere before calling this method!
 *
 * @param argId1 ID of the amino acid that may be moved if this method returns true.
 * @param argX X coordinate of the "L" location (where we may move amino acid)
 * @param argY Y coordinate of the "L" location
 */
bool SimpleHP2DRectangularAlgorithm::checkCorner(int argId, int& argX, int& argY) {
  bool ret = false; // by default we will set that corner location is not valid
  // get coordinates of the test amino acid
  int a1x = _proteinRef.get(argId).x;
  int a1y = _proteinRef.get(argId).y;
  // get coordinates of the next amino acid in the chain
  int a2x = _proteinRef.get(argId+1).x;
  int a2y = _proteinRef.get(argId+1).y;
  // coordinates of the "L" position
  int lx = argX;
  int ly = argY;
  // now calculate location of the "C" location
  int cx = lx-a2x+a1x;
  int cy = ly-a2y+a1y;
#ifdef DEBUG
      std::cout << argId << "C2:" << cx << "," << cy << std::endl;
#endif
  // get coordinates of the previous amino acid in the chain, but only if
  // we are not checking the first element in the chain.
  // test for the first case - if we have an amino acid with ID = argId-1 at
  // C-location, in that case method should return true.
  int px = 0;
  int py = 0;
  if (argId > 0) {
    px = _proteinRef.get(argId-1).x;
    py = _proteinRef.get(argId-1).y;
    if ((cx == px) && (cy == py)) {
      ret = true;
      _moveDirectly = true; // update the flag which tell us whether move can be done immediately
    } else {
      _moveDirectly = false;
    }
  } // if
  // second possible case is when there is no amino acid at C location
  if (isFree(cx, cy)) {
    ret = true;
    _cornerFree = true;
  } else {
    _cornerFree = false;
  }
  // Overwrite given coordinates for "L" location with calculated C-coordinates
  argX = cx;
  argY = cy;
  return ret;
} // checkCorner() method

/**
 * Returns a bitset in form of an unsigned char value where each bit is set to
 * true if in that direction we have a free location. In other words an amino
 * acid with index equal to "arg" can move in that direction.
 * Example: A hydrophobic amino acid and indices of bits that denote direction.
 *  3 2 1      1 - NE direction,  2 - N direction,  3 - NW direction
 *  4 H 0      0 - E direction,   4 - W direction
 *  5 6 7      5 - SW direction,  6 - S direction,  7 - SE direction.
 * NOTE: This method uses Protein's coordinate system. On the screen these directions
 * are different! Also note that available moves depend on algorithm as well.
 *
 * The implementation below follows rules described in "A complete and Effective
 * Move Set for Simplified Protein Folding" (Lesh et al., 2003). Lesh defined
 * pull moves as only diagonal ones. This simple algorithm will use the same.
 */
unsigned char SimpleHP2DRectangularAlgorithm::getAvailableMoves(unsigned int arg) {
  unsigned char ret = 0; // by default all places around the current
                         // amino acid are occupied, which is possible only
                         // in case when we have only one amino acid in the chain.
  AminoAcid& aa = _proteinRef.get(arg);
  // we have few possible cases here:
  // #1 polypeptide is empty
  if (_proteinRef.size() == 0)
    return 0;
  // #2 popypeptide contains only 1 residue
  if (_proteinRef.size() == 1)
    return 0xAA; // a single amino-acid can move diagonally, everywhere
  // #3 Give ID is id of the last amino acid - accordint to Lesh's paper it cannot be moved as well
  if (arg == _proteinRef.size()-1)
    return 0;

  // Now let's determine directions for valid "L" nodes (Lesh et al., 2003, pg189)

  // let's set initially possible moves (all diagonal moves) to 1
  //if (isFree(aa.x+1, aa.y)) bitOn(0, ret); // direction bit 0, east.
  if (isFree(aa.x+1, aa.y+1)) bitOn(1, ret); // direction bit 1, north-east.
  //if (isFree(aa.x, aa.y+1)) bitOn(2, ret); // direction bit 2, north.
  if (isFree(aa.x-1, aa.y+1)) bitOn(3, ret); // direction bit 3, north-west.
  //if (isFree(aa.x-1, aa.y)) bitOn(4, ret); // direction bit 4, west.
  if (isFree(aa.x-1, aa.y-1)) bitOn(5, ret); // direction bit 5, south-west.
  //if (isFree(aa.x, aa.y-1)) bitOn(6, ret); // direction bit 6, south.
  if (isFree(aa.x+1, aa.y-1)) bitOn(7, ret); // direction bit 7, south-east.

  // Now turn OFF those directions towards locations that are NOT adjacent
  // to the next amino acid in the chain.
  int lx = aa.x+1; // first possible L location to check - the one in the NE direction
  int ly = aa.y+1;
  if (!_proteinRef.adjacentTo(lx, ly, arg+1)) {
    bitOff(1, ret); // not adjacent, so not valid L location
  } else {
    // valid L location, but now we must check whether C location is valid as well
    if (!checkCorner(arg, lx, ly)) {
      bitOff(1, ret); // C node is not valid, so we will flag this move as invalid
    } // if
  } // else

  // check NW now
  lx = aa.x-1;
  ly = aa.y+1;
  if (!_proteinRef.adjacentTo(lx, ly, arg+1)) {
    bitOff(3, ret);
  } else {
    if (!checkCorner(arg, lx, ly)) {
      bitOff(3, ret); // C node is not valid, so we will flag this move as invalid
    } // if
  } // else

  // check SW now
  lx = aa.x-1;
  ly = aa.y-1;
  if (!_proteinRef.adjacentTo(lx, ly, arg+1)) {
    bitOff(5, ret);
  } else {
    if (!checkCorner(arg, lx, ly)) {
      bitOff(5, ret); // C node is not valid, so we will flag this move as invalid
    } // if
  } // else

  // check SW now
  lx = aa.x+1;
  ly = aa.y-1;
  if (!_proteinRef.adjacentTo(lx, ly, arg+1)) {
    bitOff(7, ret);
  } else {
    if (!checkCorner(arg, lx, ly)) {
      bitOff(7, ret); // C node is not valid, so we will flag this move as invalid
    } // if
  } // else

  // Now we should have correct value in ret, we can return it
  return ret;
} // getAvailableMoves() method

/**
 * Scans protein, and updates each amino acid in it with information about
 * available moves. This information is needed when user selects a certain
 * amino acid, to know immediately (without scanning whole protein) which
 * move can be done and update VIEW accurately.
 */
void SimpleHP2DRectangularAlgorithm::updateFreeMoves() {
  unsigned char tmp;
  for (int i=0; i<_proteinRef.size(); i++) {
    tmp = getAvailableMoves(i);
    _proteinRef.get(i).free = tmp;
  } // for
} // updateFreeMoves()

/**
 * A convenient way to initialise protein. It considers all characters in "01HPhp" set.
 * All other characters are simply ignored.
 * @param arg String (C-string) containing input sequence.
 */
void SimpleHP2DRectangularAlgorithm::initProtein(const char* arg) {
  _proteinRef.parseInput(arg);
  _proteinRef.energy(energy());
} // initProtein() method

 /**
  * gridType() returns the type of grid (2D/3D Rectangular/Triangular). For now
  * there are just 4 options.
  * @see GridType
  */
GridType SimpleHP2DRectangularAlgorithm::gridType() {
  return Rectangular2D;
}

/**
 * Overloaded is3D() method which simply returns false for this algorithm type.
 */
bool SimpleHP2DRectangularAlgorithm::is3D() {
  return false;
}

bool SimpleHP2DRectangularAlgorithm::isRectangular() {
  return true;
}

/**
 * Utility method used in case when we want to explicitly allocate a new Protein
 * and pass it to the IAlgorithm object.
 */
void SimpleHP2DRectangularAlgorithm::setProtein(Protein& arg) {
  _proteinRef = arg;
  _proteinRef.authors(_authors);
  _proteinRef.description(_description);
} // setProtein() method

/**
 * One of the crucial methods in every protein folding algorithm.
 * It calculates the energy of the current protein conformation.
 *
 * An implementation of the energy function described in the publication
 * "A replica exchange Monte Carlo algorithm for protein folding in the HP model"
 * (Thachuk et al., 2007)
 *
 * Complexity: O(n^2) where n = number of residues in the Protein.
 *
 * @param arg A reference to a Protein object that needs to be examined.
 */
double SimpleHP2DRectangularAlgorithm::energy(const Protein& arg) const {
  if (_proteinRef.size() < 2)
    return 0.0;
  double result = 0.0;
  unsigned int s = _proteinRef.size();
  for (int i = 0; i<s-1; i++) {
    for (int j=i+1; j<s; j++) {
      AminoAcid& first = _proteinRef.get(i);
      AminoAcid& second = _proteinRef.get(j);
      if ((topologicalNeighbour(first, second)) && (first.h && second.h)) {
        result = result - 1.0;
      } // if
    } // for
  } // for
  return result;
} // energy() method

/**
 * Uses method above to calculate the energy of a Protein this algorithm currently
 * works with.
 */
double SimpleHP2DRectangularAlgorithm::energy() const {
  return energy(_proteinRef);
} // energy() method

#ifdef DEBUG
#  include <iostream>
#  include <profo/Protein.h>
void SimpleHP2DRectangularAlgorithm::unittest() {
  profo::Protein p;
  p.parseInput("HHPHHPHHH");
  SimpleHP2DRectangularAlgorithm o(p);
  std::cout << "SimpleHP2DRectangularAlgorithm: " << o.energy() << std::endl;
} // unittest() method
#endif // DEBUG

// $Id: SimpleHP2DRectangularAlgorithm.cpp 44 2009-04-26 14:39:06Z dejan.lekic $
